from llama_index.llms.openrouter import OpenRouter
from llama_index.llms.openai_like import OpenAILike
from llama_index.llms.huggingface import HuggingFaceLLM
from .enum_class import LLMType

# llama_index.llms 参考文档：https://docs.llamaindex.ai/en/stable/examples/llm/


class LLMHandler:
    """
    LLM处理器
    """

    def __init__(self, llm_type: LLMType, **kwargs):
        # 初始化默认值，确保属性始终存在
        self.llm_model_name = None

        self.llm_type = llm_type
        if "llm_model_name" in kwargs:
            self.llm_model_name = kwargs["llm_model_name"]
            # 在kwargs中删除llm_model_name
            kwargs.pop("llm_model_name")
        if "api_key" in kwargs:
            self.api_key = kwargs["api_key"]
            # 在kwargs中删除api_key
            kwargs.pop("api_key")
        self.kwargs = kwargs

    def get_llm(self):
        """
        获取LLM模型，统一入口
        """
        match self.llm_type:
            case LLMType.Local_HuggingFace:
                return self.get_huggingface_llm(**self.kwargs)
            case LLMType.Online_OpenRouter:
                return self.get_openrouter_llm(self.api_key, **self.kwargs)
            case LLMType.Online_OpenAILike:
                return self.get_openailike_llm(self.api_key, **self.kwargs)
            case _:
                raise ValueError(f"不支持的LLM类型: {self.llm_type}")

    def get_openrouter_llm(self, api_key, **kwargs):
        """
        获取OpenRouter的LLM模型
        参数：
            api_key: OpenRouter的API密钥
            kwargs: 其他参数
                可选参数:
                    llm_model_name: 模型名称，默认"deepseek/deepseek-chat-v3-0324:free"
                    temperature: 温度，默认1.0
                    max_tokens: 最大tokens，默认4096
        返回：
            OpenRouter: OpenRouter的LLM模型
        """
        default_model_name = "deepseek/deepseek-chat-v3-0324:free"
        return OpenRouter(
            api_key=api_key,
            model=self.llm_model_name if self.llm_model_name else default_model_name,
            temperature=kwargs.get("temperature", 1.0),
            max_tokens=kwargs.get("max_tokens", 4096),
        )

    def get_openailike_llm(self, api_key, **kwargs):
        """
        获取兼容OpenAI格式的模型
        参数：
            api_key: OpenAI兼容格式模型的API密钥
            kwargs: 其他参数
                可选参数:
                    llm_model_name: 模型名称，默认"deepseek-chat"
                    temperature: 温度，默认1.0
                    api_base: API基础URL，默认"https://api.deepseek.com/v1"
                    max_tokens: 最大tokens，默认4096
                    is_chat_model: 是否为聊天模型，默认True
                    is_function_calling_model: 是否为函数调用模型，默认False
        返回：
            OpenAILike: 兼容OpenAI格式的模型
        """
        default_model_name = "deepseek-chat"
        return OpenAILike(
            api_key=api_key,
            model=self.llm_model_name if self.llm_model_name else default_model_name,
            temperature=kwargs.get("temperature", 1.0),
            api_base=kwargs.get("api_base", "https://api.deepseek.com/v1"),
            max_tokens=kwargs.get("max_tokens", 4096),
            is_chat_model=kwargs.get("is_chat_model", True),
            is_function_calling_model=kwargs.get("is_function_calling_model", False),
        )

    def get_huggingface_llm(self, **kwargs):
        """
        获取HuggingFace的LLM模型
        参数：
            kwargs: 其他参数
                可选参数:
                    llm_model_name: 模型名称，默认"Langboat/bloom-1b4-zh"
                    tokenizer_name: 分词器名称，默认与模型名称相同
                    device_map: 设备映射，默认"auto"，如需使用GPU，请设置为"cuda"
                      - 如果使用cuda，电脑需要安装CUDA，并安装pytorch等组件，参考如下（会自动安装适合版本的pytorch）：
                      # pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
        返回：
            HuggingFaceLLM: HuggingFace的LLM模型
        """
        default_model_name = "Langboat/bloom-1b4-zh"
        return HuggingFaceLLM(
            model_name=self.llm_model_name
            if self.llm_model_name
            else default_model_name,
            tokenizer_name=kwargs.get(
                "tokenizer_name",
                self.llm_model_name if self.llm_model_name else default_model_name,
            ),
            device_map=kwargs.get("device_map", "auto"),
        )
